import json
from redis.asyncio import Redis
import typing as T
import logging

from ...init.RedisDatabase import get_redis_session


class _ChatHistoryCacheMapper:
    redis: Redis = None
    generator: T.AsyncGenerator = None

    async def __aenter__(self):
        self.generator = get_redis_session()
        self.redis = await anext(self.generator)

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        try:
            await self.generator.aclose()
        except Exception as e:
            ...
        finally:
            self.generator = None
            self.redis = None

    async def get_cache_by_session_id(self, sessionId: int, limit: int = 10):
        async with self:
            history = await self.redis.lrange(
                f"chat_history:{sessionId}", start=-limit, end=-1
            )
            history = list(map(lambda x: json.loads(x.decode()), history))
            return history

    async def set_cache_by_session_id(
        self, sessionId: int, history: list[list[str]], ttl: int = 60 * 60 * 2
    ):
        async with self:
            pipe = self.redis.pipeline(transaction=True)
            if len(history) == 0:
                pipe.set(f"persisted:{sessionId}", len(history))
                pipe.expire(f"persisted:{sessionId}", ttl)
                await pipe.execute()
                return

            pipe.lpush(
                f"chat_history:{sessionId}",
                *list(map(lambda x: json.dumps(x, ensure_ascii=False), history)),
            )

            pipe.set(f"persisted:{sessionId}", len(history))
            pipe.expire(f"chat_history:{sessionId}", ttl)
            pipe.expire(f"persisted:{sessionId}", ttl)
            await pipe.execute()

    async def add_cache_by_session_id(
        self, sessionId: int, history: list[list[str]], ttl: int = 60 * 60 * 2
    ):
        async with self:
            pipe = self.redis.pipeline(transaction=True)
            pipe.rpush(
                f"chat_history:{sessionId}",
                *list(map(lambda x: json.dumps(x, ensure_ascii=False), history)),
            )
            pipe.expire(f"chat_history:{sessionId}", ttl)
            await pipe.execute()
            logging.debug("添加每一轮的对话聊天记录")

    async def delete_cache_by_session_id(self, sessionId: int, deleteCount: int):
        async with self:
            pipe = self.redis.pipeline(transaction=True)
            pipe.lpop(f"chat_history:{sessionId}", deleteCount)
            pipe.get(f"persisted:{sessionId}")
            result = await pipe.execute()
            logging.debug("弹出聊天记录")
            return result

    async def get_persisted(self, sessionId: int):
        async with self:
            persisted = await self.redis.get(f"persisted:{sessionId}")
            return int(persisted) if persisted else None

    async def set_persisted(self, sessionId: int, persisted: int, ttl=60 * 60 * 2):
        async with self:
            pipe = self.redis.pipeline(transaction=True)
            pipe.set(f"persisted:{sessionId}", persisted)
            pipe.expire(f"persisted:{sessionId}", ttl)
            await pipe.execute()

    async def get_cache_length_by_session_id(self, sessionId: int):
        async with self:
            length = await self.redis.llen(f"chat_history:{sessionId}")
            return length

    async def exists(self, sessionId: int) -> bool:
        async with self:
            exists = await self.redis.exists(f"chat_history:{sessionId}")
            return exists == 0

    async def keys(self, pattern: str):
        async with self:
            keys = await self.redis.keys(pattern)
            return [key.decode() for key in keys]


_chatHistoryCacheMapper: _ChatHistoryCacheMapper = _ChatHistoryCacheMapper()
